/**
 *
 * Copyright (C) 2008  Igor Vdovichenko
 *
 * This library is free software; you can redistribute it and/or
 * modify it under the terms of the GNU Lesser General Public
 * License as published by the Free Software Foundation; either
 * version 2.1 of the License, or (at your option) any later version.
 *
 * Though a sincere effort has been made to deliver a professional,
 * quality product,the library itself is distributed WITHOUT ANY WARRANTY;
 * See the GNU Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public
 * License along with this library; if not, write to the Free Software
 */
package com.vin.scriptutils.utils.ssh.trilead;

import com.trilead.ssh2.ChannelCondition;
import com.trilead.ssh2.Connection;
import com.trilead.ssh2.Session;
import com.vin.scriptutils.utils.ssh.SshTask;
import com.vin.scriptutils.utils.ssh.SshTaskExec;
import com.vin.scriptutils.utils.ssh.exception.SshTaskExecException;
import com.vin.scriptutils.utils.ssh.SshTaskResult;
import com.vin.scriptutils.utils.ssh.exception.SshAuthenticationFailedException;
import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.InterruptedIOException;
import java.io.OutputStream;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 *
 * @author igor
 */
public class TrileadSshExec implements SshTaskExec {

    private final Logger logger = LoggerFactory.getLogger(getClass());
    private final TrileadSshConnection conn;

    public TrileadSshExec(TrileadSshConnection connection) {
        this.conn = connection;
    }

    public SshTaskResult execute(SshTask task) throws SshTaskExecException, InterruptedException, IOException {

        Connection connection = conn.setConnection();
        try {

            connection.connect(null, conn.getConnectTimeout(), conn.getKeyExchangeTimeout());
            

            // if server support aes128, set this algo as default
//            final String[] ciphers =  connection.getAvailableCiphers();
//            Collection<String> ciphList = Arrays.asList(ciphers);
//            if(ciphList.contains("blowfish-ctr")){
//                String[] newciphers = new String[2];
//                newciphers[0] = new String ("aes128-ctr");
//                newciphers[1] = new String ("aes128-cbc");
//                conn.getConnection().setServer2ClientCiphers(newciphers);
//            }

            /* Authenticate */
            boolean isAuthenticated = conn.authenticate();
            

            if (!isAuthenticated) {
                throw new SshAuthenticationFailedException("Authentication failed.");
            }

            Session session = connection.openSession();
            try {
                final String command = task.getCommand();

                logger.debug("Executing: " + command);
                session.execCommand(command);

                ByteArrayOutputStream outputStreamBuffer = new ByteArrayOutputStream();
                ByteArrayOutputStream errorStreamBuffer = new ByteArrayOutputStream();
                final InputStream inputStream = task.getInputStream();

                try {
                    pumpStreams(session, inputStream, outputStreamBuffer, errorStreamBuffer);
                } catch (IOException ex) {

                    logger.error(ex.getMessage(), ex);
                } finally {
                    if (inputStream != null) {
                        inputStream.close();
                    }
                }

                return new TrileadSshTaskResult(session.getExitStatus(),
                        new ByteArrayInputStream(outputStreamBuffer.toByteArray()),
                        new ByteArrayInputStream(errorStreamBuffer.toByteArray()));
            } finally {
                session.close();
            }
        } finally {
            connection.close();
        }
    }

    private void pumpStreams(Session session, InputStream inputStream,
            OutputStream outputStream, OutputStream errorStream) throws IOException {

        OutputStream stdin = session.getStdin();
        InputStream stdout = session.getStdout();
        InputStream stderr = session.getStderr();

        byte[] buffer = new byte[8192];
        while (true) {

            if (inputStream != null) {
                int len = inputStream.read(buffer);
                if (len > 0) {
                    stdin.write(buffer, 0, len);
                    stdin.flush();
                } else {
                    stdin.close();
                    inputStream = null;
                }
            }

            if (Thread.interrupted()) {
                throw new InterruptedIOException();
            }

            if ((stdout.available() == 0) && (stderr.available() == 0)) {
                int conditions = session.waitForCondition(ChannelCondition.STDOUT_DATA | ChannelCondition.STDERR_DATA | ChannelCondition.EOF, conn.getWaitForConditionTimeout());

                if ((conditions & ChannelCondition.TIMEOUT) != 0) {
                    continue;
                }

                /* Here we do not need to check separately for CLOSED, since CLOSED implies EOF */
                if ((conditions & ChannelCondition.EOF) != 0) {
                    /* The remote side won't send us further data... */
                    if ((conditions & (ChannelCondition.STDOUT_DATA | ChannelCondition.STDERR_DATA)) == 0) {
                        /* ... and we have consumed all data in the local arrival window. */
                        break;
                    }
                }

            // if ((conditions & (ChannelCondition.STDOUT_DATA | ChannelCondition.STDERR_DATA)) == 0)
            //	throw new IllegalStateException("Unexpected condition result (" + conditions + ")");
            }

            while (stdout.available() > 0) {
                int len = stdout.read(buffer);
                if (len > 0) {
                    outputStream.write(buffer, 0, len);
                }
            }

            while (stderr.available() > 0) {
                int len = stderr.read(buffer);
                if (len > 0) {
                    errorStream.write(buffer, 0, len);
                }
            }
        }
    }
}
